package example.util;

import lombok.extern.log4j.Log4j2;
import org.apache.commons.io.IOUtils;
import org.bouncycastle.bcpg.ArmoredInputStream;
import org.bouncycastle.bcpg.ArmoredOutputStream;
import org.bouncycastle.bcpg.BCPGOutputStream;
import org.bouncycastle.bcpg.HashAlgorithmTags;
import org.bouncycastle.jce.provider.BouncyCastleProvider;
import org.bouncycastle.openpgp.*;
import org.bouncycastle.openpgp.jcajce.JcaPGPObjectFactory;
import org.bouncycastle.openpgp.operator.PublicKeyDataDecryptorFactory;
import org.bouncycastle.openpgp.operator.jcajce.*;
import org.bouncycastle.util.io.Streams;

import javax.xml.bind.DatatypeConverter;
import java.io.*;
import java.nio.charset.StandardCharsets;
import java.security.*;
import java.util.*;

/**
 * <a href="https://github.com/bcgit/bc-java/blob/main/pg/src/main/java/org/bouncycastle/openpgp/examples/PGPExampleUtil.java">Breadcrumbs bc-java</a>
 * <a href="https://itecnote.com/tecnote/java-bouncycastle-pgp-decrypt-and-verify/"></a>
 */
@Log4j2
public class PGPHelper {

    static {
        try {
            Security.addProvider(new BouncyCastleProvider());
        } catch (Exception e) {
            log.error("Security.addProvider(new BouncyCastleProvider()) err={}", e.getMessage(), e);
        }
    }

    public static PGPPrivateKey readPrivateKey(String pgpPrivateKeyAsString, String passPhrase) throws IOException, PGPException {
        PGPSecretKey pgpSecretKey = readSecretKey(IOUtils.toInputStream(pgpPrivateKeyAsString, StandardCharsets.UTF_8));
        PGPSecretKeyRingCollection pgpSecretKeyRings = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(IOUtils.toInputStream(pgpPrivateKeyAsString, StandardCharsets.UTF_8)), new JcaKeyFingerprintCalculator());
        return findSecretKey(pgpSecretKeyRings, pgpSecretKey.getKeyID(), passPhrase.toCharArray());
    }

    public static PGPPrivateKey readPrivateKey(String pgpPrivateKeyAsString, String passPhrase, long keyId) throws IOException, PGPException {
        PGPSecretKeyRingCollection pgpSecretKeyRings = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(IOUtils.toInputStream(pgpPrivateKeyAsString, StandardCharsets.UTF_8)), new JcaKeyFingerprintCalculator());
        return findSecretKey(pgpSecretKeyRings, keyId, passPhrase.toCharArray());
    }

    public static PGPPrivateKey findSecretKey(PGPSecretKeyRingCollection pgpSec, long keyID, char[] pass)
            throws PGPException {
        PGPSecretKey pgpSecKey = pgpSec.getSecretKey(keyID);

        if (pgpSecKey == null) {
            return null;
        }

        return pgpSecKey.extractPrivateKey(new JcePBESecretKeyDecryptorBuilder().setProvider("BC").build(pass));
    }


    public static PGPPublicKey readPublicKey(String pgpPublicKeyAsString) throws IOException, PGPException {
        InputStream keyIn = new ByteArrayInputStream(pgpPublicKeyAsString.getBytes());
        PGPPublicKey pubKey = readPublicKey(new ByteArrayInputStream(pgpPublicKeyAsString.getBytes()));
        keyIn.close();
        return pubKey;
    }

    public static PGPPublicKey readPublicKey2(String pgpPublicKeyAsString) throws IOException, PGPException {
        InputStream keyIn = new ByteArrayInputStream(pgpPublicKeyAsString.getBytes());
        PGPPublicKey pubKey = readPublicKey2(new ByteArrayInputStream(pgpPublicKeyAsString.getBytes()));
        keyIn.close();
        return pubKey;
    }

    /**
     * A simple routine that opens a key ring file and loads the first available key
     * suitable for encryption.
     *
     * @param input data stream containing the public key data
     * @return the first public key found.
     */
    public static PGPPublicKey readPublicKey2(InputStream input) throws IOException, PGPException {
        PGPPublicKeyRingCollection pgpPub = new PGPPublicKeyRingCollection(PGPUtil.getDecoderStream(input),
                new JcaKeyFingerprintCalculator());

        //
        // we just loop through the collection till we find a key suitable for
        // encryption, in the real
        // world you would probably want to be a bit smarter about this.
        //

        Iterator<PGPPublicKeyRing> keyRingIter = pgpPub.getKeyRings();
        while (keyRingIter.hasNext()) {
            PGPPublicKeyRing keyRing = keyRingIter.next();

            Iterator<PGPPublicKey> keyIter = keyRing.getPublicKeys();
            while (keyIter.hasNext()) {
                PGPPublicKey key = keyIter.next();
                if (key.isMasterKey()) {
                    return key;
                }
            }
        }

        throw new IllegalArgumentException("Can't find encryption key in key ring.");
    }

    /**
     * A simple routine that opens a key ring file and loads the first available key
     * suitable for encryption.
     *
     * @param input data stream containing the public key data
     * @return the first public key found.
     */
    public static PGPPublicKey readPublicKey(InputStream input) throws IOException, PGPException {
        PGPPublicKeyRingCollection pgpPub = new PGPPublicKeyRingCollection(PGPUtil.getDecoderStream(input),
                new JcaKeyFingerprintCalculator());

        //
        // we just loop through the collection till we find a key suitable for
        // encryption, in the real
        // world you would probably want to be a bit smarter about this.
        //

        Iterator<PGPPublicKeyRing> keyRingIter = pgpPub.getKeyRings();
        while (keyRingIter.hasNext()) {
            PGPPublicKeyRing keyRing = keyRingIter.next();

            Iterator<PGPPublicKey> keyIter = keyRing.getPublicKeys();
            while (keyIter.hasNext()) {
                PGPPublicKey key = keyIter.next();
                if (key.isEncryptionKey() && !key.isMasterKey()) {
                    return key;
                }
            }
        }

        throw new IllegalArgumentException("Can't find encryption key in key ring.");
    }

    public static String encrpytRequestPayload(String requestPayloadPlainString, PGPPublicKey publicKey,
                                               boolean armor, boolean withIntegrityCheck, boolean isMasterKey) throws IOException, PGPException {
        ByteArrayOutputStream out = new ByteArrayOutputStream();
        OutputStream opOut;
        if (armor) {
            opOut = new ArmoredOutputStream(out);
        } else {
            opOut = out;
        }

        ByteArrayOutputStream byteArray = new ByteArrayOutputStream();

        InputStream inputStream = IOUtils.toInputStream(requestPayloadPlainString, StandardCharsets.UTF_8);

        PGPLiteralDataGenerator literalDataGenerator = new PGPLiteralDataGenerator();
        OutputStream outputStream = literalDataGenerator.open(byteArray, PGPLiteralData.BINARY, PGPLiteralData.CONSOLE,
                inputStream.available(), new Date());
        Streams.pipeAll(inputStream, outputStream);
        outputStream.close();

        JcePGPDataEncryptorBuilder dataEncryptorBuilder = new JcePGPDataEncryptorBuilder(PGPEncryptedData.AES_256)
                .setWithIntegrityPacket(withIntegrityCheck).setSecureRandom(new SecureRandom()).setProvider("BC");

        PGPEncryptedDataGenerator encryptedDataGenerator = new PGPEncryptedDataGenerator(dataEncryptorBuilder);
        PGPPublicKey encryptionKey = null;

        for (PGPPublicKey key : Collections.singletonList(publicKey)) {
            if (key.isEncryptionKey() && (!key.isMasterKey() || isMasterKey)) {
                encryptionKey = key;
                break;
            }
        }
        if (encryptionKey == null) {
            throw new IllegalStateException("matching encrypted data not found");
        }
        JcePublicKeyKeyEncryptionMethodGenerator encryptionMethodGenerator = new JcePublicKeyKeyEncryptionMethodGenerator(
                encryptionKey).setProvider("BC").setSecureRandom(new SecureRandom());

        encryptedDataGenerator.addMethod(encryptionMethodGenerator);

        byte[] bytes = byteArray.toByteArray();
        try (OutputStream dbOut = encryptedDataGenerator.open(opOut, bytes.length)) {
            dbOut.write(bytes);
            dbOut.close();
            opOut.close();
            return out.toString();
        }
    }

    // Helper method for getting message digest
    static String generateMessageDigest(String message, String hashingAlgorithm)
            throws NoSuchAlgorithmException {
        return DatatypeConverter.printHexBinary(MessageDigest.getInstance(hashingAlgorithm).digest(message.getBytes()))
                .toLowerCase();
    }


    /**
     * A simple routine that opens a key ring file and loads the first available key
     * suitable for signature generation.
     *
     * @param input stream to read the secret key ring collection from.
     * @return a secret key.
     * @throws IOException  on a problem with using the input stream.
     * @throws PGPException if there is an issue parsing the input stream.
     */
    public static PGPSecretKey readSecretKey2(InputStream input) throws IOException, PGPException {
        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(input),
                new JcaKeyFingerprintCalculator());

        //
        // we just loop through the collection till we find a key suitable for
        // encryption, in the real
        // world you would probably want to be a bit smarter about this.
        //

        Iterator<PGPSecretKeyRing> keyRingIter = pgpSec.getKeyRings();
        while (keyRingIter.hasNext()) {
            PGPSecretKeyRing keyRing = keyRingIter.next();

            Iterator<PGPSecretKey> keyIter = keyRing.getSecretKeys();
            while (keyIter.hasNext()) {
                PGPSecretKey key = keyIter.next();

                if (key.isMasterKey()) {
                    return key;
                }
            }
        }

        throw new IllegalArgumentException("Can't find signing key in key ring.");
    }

    /**
     * A simple routine that opens a key ring file and loads the first available key
     * suitable for signature generation.
     *
     * @param input stream to read the secret key ring collection from.
     * @return a secret key.
     * @throws IOException  on a problem with using the input stream.
     * @throws PGPException if there is an issue parsing the input stream.
     */
    public static PGPSecretKey readSecretKey(InputStream input) throws IOException, PGPException {
        PGPSecretKeyRingCollection pgpSec = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(input),
                new JcaKeyFingerprintCalculator());

        //
        // we just loop through the collection till we find a key suitable for
        // encryption, in the real
        // world you would probably want to be a bit smarter about this.
        //

        Iterator<PGPSecretKeyRing> keyRingIter = pgpSec.getKeyRings();
        while (keyRingIter.hasNext()) {
            PGPSecretKeyRing keyRing = keyRingIter.next();

            Iterator<PGPSecretKey> keyIter = keyRing.getSecretKeys();
            while (keyIter.hasNext()) {
                PGPSecretKey key = keyIter.next();

                if (key.isSigningKey() && !key.isMasterKey()) {
                    return key;
                }
            }
        }

        throw new IllegalArgumentException("Can't find signing key in key ring.");
    }

    public static void decryptVerifyMessage(InputStream in, OutputStream out, String  privateKeyString, String passPhrase, PGPPublicKey publicKey) throws Exception {
        in = new ArmoredInputStream(in);

        JcaKeyFingerprintCalculator keyFingerPrintCalculator = new JcaKeyFingerprintCalculator();

        PGPObjectFactory pgpF = new PGPObjectFactory(in, keyFingerPrintCalculator);
        PGPEncryptedDataList enc = (PGPEncryptedDataList) pgpF.nextObject();
        PGPPublicKeyEncryptedData encryptedData = (PGPPublicKeyEncryptedData) enc.getEncryptedDataObjects().next();
        long privateKeyId = encryptedData.getKeyID();

        // 协商获取私钥
        PGPPrivateKey privateKey = readPrivateKey(privateKeyString, passPhrase, privateKeyId);


        PGPObjectFactory plainFact = new PGPObjectFactory(encryptedData.getDataStream(new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC").build(privateKey)), keyFingerPrintCalculator);

        Object message = null;

        PGPOnePassSignatureList onePassSignatureList = null;
        PGPSignatureList signatureList = null;
        PGPCompressedData compressedData = null;

        message = plainFact.nextObject();
        ByteArrayOutputStream actualOutput = new ByteArrayOutputStream();

        while (message != null) {
            System.out.println(message);
            if (message instanceof PGPCompressedData) {
                compressedData = (PGPCompressedData) message;
                plainFact = new PGPObjectFactory(compressedData.getDataStream(), keyFingerPrintCalculator);
                message = plainFact.nextObject();
                System.out.println(message.toString());
            }

            if (message instanceof PGPLiteralData) {
                Streams.pipeAll(((PGPLiteralData) message).getInputStream(), actualOutput);
            } else if (message instanceof PGPOnePassSignatureList) {
                onePassSignatureList = (PGPOnePassSignatureList) message;
            } else if (message instanceof PGPSignatureList) {
                signatureList = (PGPSignatureList) message;
            } else {
                throw new PGPException("message unknown message type.");
            }
            message = plainFact.nextObject();
        }
        actualOutput.close();
        byte[] output = actualOutput.toByteArray();
        if (onePassSignatureList == null || signatureList == null) {
            throw new PGPException("Poor PGP. Signatures not found.");
        } else {

            for (int i = 0; i < onePassSignatureList.size(); i++) {
                PGPOnePassSignature ops = onePassSignatureList.get(0);
                System.out.println("verifier : " + ops.getKeyID());
                if (publicKey != null) {
                    ops.init(new JcaPGPContentVerifierBuilderProvider().setProvider("BC"), publicKey);
                    ops.update(output);
                    PGPSignature signature = signatureList.get(i);
                    if (ops.verify(signature)) {
                        Iterator<?> userIds = publicKey.getUserIDs();
                        while (userIds.hasNext()) {
                            String userId = (String) userIds.next();
                            System.out.println("Signed by " + userId);
                        }
                        System.out.println("Signature verified");
                    } else {
                        throw new SignatureException("Signature verification failed");
                    }
                }
            }

        }

        out.write(output);
        out.flush();
        out.close();
    }

    public static String decryptResponseMessage(String pgpPrivateKeyAsString, String passPhrase, String responseMessageEncryptedBase64String)
            throws PGPException, IOException {

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        OutputStream out = new ArmoredOutputStream(outputStream);

        try (InputStream pgpEncryptedData = PGPUtil.getDecoderStream(IOUtils.toInputStream(responseMessageEncryptedBase64String, StandardCharsets.UTF_8))) {


            JcaKeyFingerprintCalculator keyFingerPrintCalculator = new JcaKeyFingerprintCalculator();

            // 解码base64
            String epicStrNew = new String(Base64.getDecoder().decode(responseMessageEncryptedBase64String));


            BufferedInputStream dataStream = new BufferedInputStream(IOUtils.toInputStream(epicStrNew));

            InputStream decoderStream = PGPUtil.getDecoderStream(dataStream);

            PGPObjectFactory pgpObjectFactory = new PGPObjectFactory(decoderStream, keyFingerPrintCalculator);

            PGPEncryptedDataList pgpEncryptedDataList = (PGPEncryptedDataList) pgpObjectFactory.nextObject();

            // find the matching public key encrypted data packet.
            PGPPublicKeyEncryptedData pgpPublicKeyEncryptedData = null;



            PGPPrivateKey privateKey = null;


            // build decryptor factory
            PublicKeyDataDecryptorFactory dataDecryptorFactory = null;
            for (Object pgpEncData : pgpEncryptedDataList) {
                PGPPublicKeyEncryptedData pkEnc = (PGPPublicKeyEncryptedData) pgpEncData;
                //for (PGPPrivateKey privateKey : privateKeys) {
                if (pkEnc.getKeyID() == privateKey.getKeyID()) {
                    pgpPublicKeyEncryptedData = pkEnc;
                    dataDecryptorFactory = new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC")
                            .build(privateKey);
                    break;
                }
                //}
            }

            PGPSecretKey pgpSecretKey = readSecretKey(IOUtils.toInputStream(pgpPrivateKeyAsString, StandardCharsets.UTF_8));
            PGPSecretKeyRingCollection pgpSecretKeyRings = new PGPSecretKeyRingCollection(PGPUtil.getDecoderStream(IOUtils.toInputStream(pgpPrivateKeyAsString, StandardCharsets.UTF_8)), new JcaKeyFingerprintCalculator());

            if (pgpPublicKeyEncryptedData == null || dataDecryptorFactory == null) {
                throw new IllegalStateException("matching encrypted data not found");
            }


            InputStream pgpPrivateKeyInputStream = pgpPublicKeyEncryptedData.getDataStream(new JcePublicKeyDataDecryptorFactoryBuilder().setProvider("BC").build(privateKey));
            JcaPGPObjectFactory pgpPrivateKey_jcaPGPObjectFactory = new JcaPGPObjectFactory(pgpPrivateKeyInputStream);
            Object message = pgpPrivateKey_jcaPGPObjectFactory.nextObject();

            if (message instanceof PGPCompressedData) {
                PGPCompressedData pgpCompressedData = (PGPCompressedData) message;
                JcaPGPObjectFactory pgpFact = new JcaPGPObjectFactory(pgpCompressedData.getDataStream());

                message = pgpFact.nextObject();
            }

            if (message instanceof PGPLiteralData) {
                PGPLiteralData pgpLiteralData = (PGPLiteralData) message;
                InputStream unc = pgpLiteralData.getInputStream();
                Streams.pipeAll(unc, out);

            } else if (message instanceof PGPOnePassSignatureList) {
                throw new PGPException("encrypted message contains a signed message - not literal data.");
            } else {
                throw new PGPException("message is not a simple encrypted file - type unknown.");
            }

//            log.info(outputStream.toString());

            if (pgpPublicKeyEncryptedData.isIntegrityProtected()) {
                if (!pgpPublicKeyEncryptedData.verify()) {
                    log.error("message failed integrity check");
                } else {
                    log.info("message integrity check passed");
                }
            } else {
                log.error("no message integrity check");
            }
        } catch (PGPException e) {
            log.error(e.getMessage(), e);
            if (e.getUnderlyingException() != null) {
                log.error(e.getUnderlyingException().getMessage(), e.getUnderlyingException());
            }
        }

        return new String(Base64.getEncoder().encode(outputStream.toByteArray()), StandardCharsets.UTF_8);
    }

    public static byte[] signatureMessage(PGPPrivateKey privateKey, String responseMessageEncryptedBase64String, boolean armor)
            throws PGPException, IOException {

        PGPSignatureGenerator sGen = new PGPSignatureGenerator(new JcaPGPContentSignerBuilder(privateKey.getPublicKeyPacket().getAlgorithm(), HashAlgorithmTags.SHA1).setProvider(new BouncyCastleProvider()));
        sGen.init(PGPSignature.BINARY_DOCUMENT, privateKey);

        ByteArrayOutputStream byteOut = new ByteArrayOutputStream();
        ArmoredOutputStream aOut = new ArmoredOutputStream(byteOut);

        BCPGOutputStream bOut = new BCPGOutputStream(byteOut);
        InputStream fIn = IOUtils.toInputStream(responseMessageEncryptedBase64String, StandardCharsets.UTF_8);

        int ch;
        while ((ch = fIn.read()) >= 0) {
            sGen.update((byte) ch);
        }

        aOut.endClearText();
        fIn.close();
        sGen.generate().encode(bOut);
        if (armor) {
            aOut.close();
        }

        return byteOut.toByteArray();
    }

    private static final Provider bcp = new BouncyCastleProvider();
    private static final int bufferSize = 4096;

    public static String encryptAndSign2base64(String msg, PGPPublicKey publicKey, PGPPrivateKey privateKey) throws PGPException, IOException {

        ByteArrayOutputStream outputStream = new ByteArrayOutputStream();
        OutputStream out = new ArmoredOutputStream(outputStream);

        PGPEncryptedDataGenerator encryptedDataGenerator = new PGPEncryptedDataGenerator(new JcePGPDataEncryptorBuilder(PGPEncryptedData.AES_256)
                .setWithIntegrityPacket(true)
                .setSecureRandom(new SecureRandom())
                .setProvider(bcp));

        encryptedDataGenerator.addMethod(new JcePublicKeyKeyEncryptionMethodGenerator(publicKey).setProvider(bcp));
        OutputStream encryptedOut = encryptedDataGenerator.open(out, new byte[bufferSize]);

        PGPCompressedDataGenerator compressedDataGenerator = new PGPCompressedDataGenerator(PGPCompressedData.ZIP);
        OutputStream compressedData = compressedDataGenerator.open(encryptedOut);

        PGPSignatureGenerator signatureGenerator = new PGPSignatureGenerator(new JcaPGPContentSignerBuilder(privateKey.getPublicKeyPacket().getAlgorithm(), PGPUtil.SHA1).setProvider("BC"));
        signatureGenerator.init(PGPSignature.BINARY_DOCUMENT, privateKey);
        Iterator<String> it = publicKey.getUserIDs();
        if (it.hasNext()) {
            PGPSignatureSubpacketGenerator spGen = new PGPSignatureSubpacketGenerator();
            spGen.addSignerUserID(false, it.next());
            signatureGenerator.setHashedSubpackets(spGen.generate());
        }
        signatureGenerator.generateOnePassVersion(false).encode(compressedData);

        PGPLiteralDataGenerator literalDataGenerator = new PGPLiteralDataGenerator();
        OutputStream lOut = literalDataGenerator.open(compressedData, PGPLiteralData.BINARY, "Sample-Data", new Date(), new byte[bufferSize]);
        byte[] data = msg.getBytes(StandardCharsets.UTF_8);
        lOut.write(data);

        signatureGenerator.update(data);
        lOut.close();
        literalDataGenerator.close();
        signatureGenerator.generate().encode(compressedData);

        compressedDataGenerator.close();
        compressedData.close();
        encryptedOut.close();
        encryptedDataGenerator.close();
        out.close();

        return new String(Base64.getEncoder().encode(outputStream.toByteArray()), StandardCharsets.UTF_8);
    }

}
